DuckDBとAWS Lambda PythonでS3にSQLを投げる環境をAWS SAMで構築してみた

quiver

2024.12.31

本記事では、Amazon S3上のCSVファイルに対して、OLAPが得意なDBMSであるDuckDB経由でコンテナ版AWS Lambda PythonからSQLを投げる環境をAWS SAMで構築する方法を紹介します。

S3上の大きすぎないCSV/Parquet/IcebergファイルをLambdaでシュッと同期的に処理したい場合に便利です。

特に、以下の3点は重要です。

Lambdaは /tmp 以下のみ書き込み可能
Lambdaは環境変数 $HOME が定義されておらず、DuckDB は同変数が定義されていることを前提とする
Lambdaの最大メモリは10GB

やってみた

S3にCSVファイルをアップロード

CSVファイルをS3にアップロードし、S3 URI(s3://YOUR-BUCKET/FILENAME.csv)を控えます。

AWS SAMテンプレートを取得

次の GitHub レポジトリから、AWS SAMテンプレートを取得します

$ git clone https://github.com/quiver/duckdb-aws-lambda-python-sam-template.git
$ cd duckdb-aws-lambda-python-sam-template

今回の構成では、問い合わせ先のS3 オブジェクトをAWS Lambdaの環境変数 S3URI で管理しています。
template.yaml の次の S3URI 変数を、書き換えます。

      Environment:
        Variables:
          HOME: /tmp
          S3URI: s3://YOUR-BUCKET/FILENAME.csv

AWS SAMでデプロイ

AWS SAM CLI をインストールし、

$ sam build
$ sam deploy --guided

Lambda関数を実行

AWSコンソール、あるいは、コンソールから、デプロイされたAWS Lambda関数を実行しましょう。

$ aws lambda invoke \
 --function-name arn:aws:lambda:ap-northeast-1:123456789012:function:DuckDBLambdaFunction \
 --payload '{}' \
 response.json
{
 "StatusCode": 200,
 "ExecutedVersion": "$LATEST"
}

$ cat response.json

... SQL 実行結果 ...

DuckDB 操作の修正

アプリケーションコードは src/app.py にあります。

CSV ではなく Parquet からの呼び出しやSQLを修正する場合は、このコードを修正してください。

import os
import duckdb

S3URI = os.environ.get('S3URI')

con = duckdb.connect()
con.execute("""
INSTALL httpfs;
LOAD httpfs;

CREATE SECRET (
      TYPE S3,
      PROVIDER CREDENTIAL_CHAIN
);
""")

def lambda_handler(event, context):
    res = con.sql(f"SELECT * FROM read_csv('{S3URI}') LIMIT 3")
    print(res)
    return str(res)

DuckDB x Lambda関数の覚書

DuckDBをLambda関数から呼び出す上で、少しハマったため、備忘録を残しておきます。

権限管理

Lambda関数がS3オブジェクトを参照できるように、Lambda関数にはS3の参照権限を付与してます。template.yaml では以下の箇所です。

Resources:
  DuckDBLambdaFunction:
    Type: AWS::Serverless::Function
    Properties:
      ...
      Policies:
        - AmazonS3ReadOnlyAccess

S3インターフェース対応

DuckDB が S3 オブジェクトを Assume Roleしてアクセスするには、 httpfs エクステンションの有効化と SECRET の定義が必要です。

src/app.py の以下の処理が対応しています。

import duckdb

...

con = duckdb.connect()
con.execute("""
INSTALL httpfs;
LOAD httpfs;

CREATE SECRET (
      TYPE S3,
      PROVIDER CREDENTIAL_CHAIN
);
""")

DuckDB エクステンションのロードでは、DuckDBはデフォルトでは $HOME 以下にファイル書き込みしようとします。

一方で、Lambdaは /tmp ディレクトリにしか書き込みできず、ワーキングディレクトリは /var/task です。さらには、 Lambda関数では $HOME 環境変数は設定されていません。そのため、Lambda関数の環境変数で $HOME を /tmp に設定しています。DuckDBセッションレベルで SET home_directory='/tmp' というように定義することも可能です。

エラー再現手順

$ HOME='' duckdb

D INSTALL httpfs;
IO Error: Can't find the home directory at ''
Specify a home directory using the SET home_directory='/path/to/dir' option.

...

D CREATE SECRET (
      TYPE S3,
      PROVIDER CREDENTIAL_CHAIN
  );
Extension Autoloading Error: An error occurred while trying to automatically install the required extension 'aws':
Can't find the home directory at ''
Specify a home directory using the SET home_directory='/path/to/dir' option.